#include <user.h>

extern void __context_switch(uint32_t *from_esp, uint32_t *from_pc, uint32_t to_esp, uint32_t to_pc);
extern void __context_init(void);

#define IO_BUFFER_INITIAL_CAP       1024

#define FIBER_STATUS_UNINIT          1
#define FIBER_STATUS_RUNNABLE_WEAK   2
#define FIBER_STATUS_RUNNABLE_STRONG 3
#define FIBER_STATUS_RUNNABLE_LOW    4
#define FIBER_STATUS_WAIT            5
#define FIBER_STATUS_IO              6
#define FIBER_STATUS_ZOMBIE          7

static
void _fiber_end(fiber_t fiber)
{
	 while (1)
	 {
		  upriv->in_sched = 1;
		  upriv->fiber_cur->status = FIBER_STATUS_ZOMBIE;
		  fiber_yield();
	 }
}

static
char *_fiber_stack_fill(char *stk_top, fiber_t fiber, fiber_entry_f entry, void *arg)
{
	 stk_top -= sizeof(uintptr_t);
	 *(uintptr_t *)stk_top = (uintptr_t)fiber;

	 stk_top -= sizeof(uintptr_t);
	 *(uintptr_t *)stk_top = 0;

	 stk_top -= sizeof(uintptr_t);
	 *(uintptr_t *)stk_top = (uintptr_t)arg;

	 stk_top -= sizeof(uintptr_t);
	 *(uintptr_t *)stk_top = (uintptr_t)_fiber_end;

	 stk_top -= sizeof(uintptr_t);
	 *(uintptr_t *)stk_top = (uintptr_t)entry;

	 return stk_top;
}


#define FIBER_STACK_SIZE      4096

static void
_fiber_switch(fiber_t fiber)
{
	 fiber_t old = upriv->fiber_cur;
	 upriv->fiber_cur = fiber;
	 
	 if (old == NULL)
	 {
		  uint32_t __unused;
		  /* Initial switch */
		  __context_switch(&__unused, &__unused, fiber->ctx.stk_ptr, fiber->ctx.pc);
	 }
	 else
	 {
		  __context_switch(&old->ctx.stk_ptr, &old->ctx.pc, fiber->ctx.stk_ptr, fiber->ctx.pc);
	 }
}


static void
fiber_kmsg_dispatcher(void **buf, uint32_t cap, uint32_t head, uint32_t tail)
{
	 if (upriv->in_sched == 0)
	 {
		  while (head != tail)
		  {
			   fiber_t f = buf[head];

			   if (f->upriv != upriv)
					kprintf("ASSERTION FAILED, RACING CONDITION!!!\n");
			   head = (head + 1) % cap;

			   if (f->status == FIBER_STATUS_RUNNABLE_WEAK)
					f->status = FIBER_STATUS_RUNNABLE_STRONG;
			   
			   else if (f->status != FIBER_STATUS_RUNNABLE_STRONG)
			   {
					f->status = FIBER_STATUS_RUNNABLE_WEAK;
					f->sched_node.next = &upriv->sched_node;
					f->sched_node.prev = f->sched_node.next->prev;
					
					f->sched_node.next->prev = &f->sched_node;
					f->sched_node.prev->next = &f->sched_node;
			   }
		  }

		  upriv->msg_incoming = 0;
		  kmsg_return(tail, 0);
	 }
	 else
	 {
		  upriv->msg_incoming = 1;
		  kmsg_return(head, 1);
	 }
}

int
fiber_init(fiber_entry_f entry, void *arg)
{
	 fiber_t f = (fiber_t)alloc(sizeof(struct fiber_s));
	 if (f == NULL) return -1;

	 uintptr_t stk = (uintptr_t)balloc(FIBER_STACK_SIZE);
	 if (stk == 0)
	 {
		  free(f);
		  return -1;
	 }

	 f->ctx.stk_ptr = (uintptr_t)_fiber_stack_fill(
		  (char *)(f->ctx.stk_top = stk + FIBER_STACK_SIZE),
		  f, entry, arg);
	 f->ctx.pc      = (uintptr_t)__context_init;
	 f->status      = FIBER_STATUS_RUNNABLE_WEAK;
	 f->uevent_touched = 0;

	 struct uthread_priv_s *priv = (struct uthread_priv_s *)alloc(sizeof(struct uthread_priv_s));
	 f->upriv = priv;

	 if (priv == NULL)
	 {
		  free(f); bfree((void *)stk);
		  return -1;
	 }

	 uint32_t *io_buffer = (uint32_t *)balloc(IO_BUFFER_INITIAL_CAP * sizeof(uint32_t));

	 if (io_buffer == NULL)
	 {
		  free(priv); free(f); bfree((void *)stk);
		  return -1;
	 }
	 
	 priv->fiber_cur = NULL;
	 priv->sched_node.next = priv->sched_node.prev = &priv->sched_node;
	 priv->low_sched_node.next = priv->low_sched_node.prev = &priv->low_sched_node;
	 f->sched_node.next = f->sched_node.prev = &f->sched_node;

	 priv->in_sched     = 0;
	 priv->msg_incoming = 0;

	 priv->io_count   = 0;
	 priv->io_buf_cap = IO_BUFFER_INITIAL_CAP;
	 priv->io_buf     = io_buffer;

	 upriv_set(priv);
	 kmsg_init(fiber_kmsg_dispatcher);
	 
	 _fiber_switch(f);
	 return 0;
}

fiber_t
fiber_create(fiber_entry_f entry, void *arg)
{
	 fiber_t f = (fiber_t)alloc(sizeof(struct fiber_s));
	 if (f == NULL) return NULL;

	 uintptr_t stk = (uintptr_t)balloc(FIBER_STACK_SIZE);
	 if (stk == 0)
	 {
		  free(f);
		  return NULL;
	 }

	 f->ctx.stk_ptr = (uintptr_t)_fiber_stack_fill(
		  (char *)(f->ctx.stk_top = stk + FIBER_STACK_SIZE),
		  f, entry, arg);
	 f->ctx.pc      = (uintptr_t)__context_init;
	 f->status      = FIBER_STATUS_RUNNABLE_WEAK;
	 f->uevent_touched = 0;
	 f->upriv = upriv;

	 upriv->in_sched = 1;
	 
	 f->sched_node.next = &upriv->sched_node;
	 f->sched_node.prev = f->sched_node.next->prev;

	 f->sched_node.next->prev = &f->sched_node;
	 f->sched_node.prev->next = &f->sched_node;

	 upriv->in_sched = 0;
	 if (upriv->msg_incoming)
		  kmsg_clean_busy();

	 return f;
}

fiber_t
fiber_self(void)
{
	 return upriv->fiber_cur;
}


static void
print_backtrace(void)
{
	 uint32_t ebp, eip;
	 __asm__ __volatile__("movl %%ebp, %0" : "=r" (ebp));
	 if (ebp == 0)
	 {
		  kprintf("Stack backtrace is empty(optimized?)\n");
		  return;
	 }

     kprintf("Stack backtrace:\n");
     /* kprintf("  ebp: %p\n", ebp); */

	 eip = *((uint32_t *)(ebp) + 1);
	 ebp = *(uint32_t *)(ebp);

     while (1)
     {
		  /* break if the ptr down to user mode */
		  if (ebp == 0)
		  {
			   break;
		  }
          
		  kprintf("  ebp: %p  eip: %p\n", ebp, eip);

		  eip = *((uint32_t *)(ebp) + 1);
		  ebp = *(uint32_t *)(ebp);
     }
}

static void
_fiber_thread_entry(void)
{
	 kmsg_init(fiber_kmsg_dispatcher);
}

int
fiber_thread_create(fiber_entry_f entry, void *arg)
{
	 fiber_t f = (fiber_t)alloc(sizeof(struct fiber_s));
	 if (f == NULL) return -1;

	 uintptr_t stk = (uintptr_t)balloc(FIBER_STACK_SIZE);
	 if (stk == 0)
	 {
		  free(f);
		  return -1;
	 }

	 f->ctx.stk_ptr = (uintptr_t)_fiber_stack_fill(
		  (char *)(f->ctx.stk_top = stk + FIBER_STACK_SIZE),
		  f, entry, arg);
	 f->ctx.pc      = 0;
	 f->status      = FIBER_STATUS_RUNNABLE_WEAK;
	 f->uevent_touched = 0;

	 struct uthread_priv_s *priv = (struct uthread_priv_s *)alloc(sizeof(struct uthread_priv_s));
	 f->upriv = priv;

	 if (priv == NULL)
	 {
		  free(f); bfree((void *)stk);
		  return -1;
	 }

	 uint32_t *io_buffer = (uint32_t *)balloc(IO_BUFFER_INITIAL_CAP * sizeof(uint32_t));

	 if (io_buffer == NULL)
	 {
		  free(priv); free(f); bfree((void *)stk);
		  return -1;
	 }
	 
	 priv->fiber_cur = f;
	 priv->sched_node.next = priv->sched_node.prev = &priv->sched_node;
	 priv->low_sched_node.next = priv->low_sched_node.prev = &priv->low_sched_node;
	 f->sched_node.next = f->sched_node.prev = &f->sched_node;

	 priv->in_sched     = 0;
	 priv->msg_incoming = 0;
	 
	 priv->io_count   = 0;
	 priv->io_buf_cap = IO_BUFFER_INITIAL_CAP;
	 priv->io_buf     = io_buffer;

	 /* TODO: STACK MAGIC! */
	 f->ctx.stk_ptr -= sizeof(uintptr_t);
	 *(uintptr_t *)(f->ctx.stk_ptr) = (uintptr_t)_fiber_thread_entry;
	 
	 uthread_create("__uproc", (uintptr_t)__context_init, (uintptr_t)priv, f->ctx.stk_ptr);
	 
	 return 0;
}

void
fiber_yield_io(void)
{
	 /* in_sched would be fixed in fiber_yield */
	 upriv->in_sched = 1;
	 upriv->fiber_cur->status = FIBER_STATUS_IO; 
	 fiber_yield();
}

void
fiber_yield_low(void)
{
	 /* in_sched would be fixed in fiber_yield */
	 upriv->in_sched = 1;
	 upriv->fiber_cur->status = FIBER_STATUS_RUNNABLE_LOW;
	 fiber_yield();
}

void
fiber_wait_pretend(void)
{
	 upriv->fiber_cur->status = FIBER_STATUS_RUNNABLE_WEAK;
}

void
fiber_wait_try(void)
{
	 /* in_sched would be fixed in fiber_yield */
	 upriv->in_sched = 1;
	 if (upriv->fiber_cur->status == FIBER_STATUS_RUNNABLE_WEAK)
	 {
		  upriv->fiber_cur->status = FIBER_STATUS_WAIT; 
		  fiber_yield();
	 }
}

void
fiber_notify(fiber_t f)
{
	 if (f->upriv == upriv)
	 {
		  upriv->in_sched = 1;

		  if (f->status == FIBER_STATUS_RUNNABLE_WEAK)
			   f->status = FIBER_STATUS_RUNNABLE_STRONG;
			   
		  else if (f->status != FIBER_STATUS_RUNNABLE_STRONG)
		  {
			   f->status = FIBER_STATUS_RUNNABLE_WEAK;
			   f->sched_node.next = &upriv->sched_node;
			   f->sched_node.prev = f->sched_node.next->prev;
			   
			   f->sched_node.next->prev = &f->sched_node;
			   f->sched_node.prev->next = &f->sched_node;
		  }
		  
		  upriv->in_sched = 0;
		  if (upriv->msg_incoming)
			   kmsg_clean_busy();
	 }
	 else fio_notify(f);
}

void
fiber_process_io_buffer(void)
{
	 uint32_t c = process_io_buffer(upriv->io_buf, upriv->io_count);

	 upriv->in_sched = 1;
	 upriv->io_count = 0;
	 while (c > 0)
	 {
		  -- c;

		  fiber_t f = (fiber_t)upriv->io_buf[c];

		  f->status = FIBER_STATUS_RUNNABLE_WEAK;

		  f->sched_node.next = &upriv->sched_node;
		  f->sched_node.prev = f->sched_node.next->prev;

		  f->sched_node.next->prev = &f->sched_node;
		  f->sched_node.prev->next = &f->sched_node;
	 }

	 upriv->in_sched = 0;
	 if (upriv->msg_incoming)
		  kmsg_clean_busy();
}

void
fiber_yield(void)
{
	 upriv->in_sched = 1;
	 
	 fiber_t f = upriv->fiber_cur;

	 if (f->status == FIBER_STATUS_RUNNABLE_WEAK ||
		 f->status == FIBER_STATUS_RUNNABLE_STRONG)
	 {
		  f->sched_node.next = &upriv->sched_node;
		  f->sched_node.prev = f->sched_node.next->prev;

		  f->sched_node.next->prev = &f->sched_node;
		  f->sched_node.prev->next = &f->sched_node;
	 }
	 else if (f->status == FIBER_STATUS_RUNNABLE_LOW)
	 {
		  f->status = FIBER_STATUS_RUNNABLE_WEAK;
		  
		  f->sched_node.next = &upriv->low_sched_node;
		  f->sched_node.prev = f->sched_node.next->prev;
		  
		  f->sched_node.next->prev = &f->sched_node;
		  f->sched_node.prev->next = &f->sched_node;
	 }

	 while (1)
	 {
		  struct sched_node_s *s = upriv->sched_node.next;
		  if (s == &upriv->sched_node)
		  {
			   if (upriv->low_sched_node.next == &upriv->low_sched_node)
			   {
					/* XXX: Why is the idle loop here? */
					upriv->in_sched = 0;
					if (upriv->msg_incoming)
						 kmsg_clean_busy();
					
					fiber_process_io_buffer();
					upriv->in_sched = 1;
					continue;
			   }
			   else
			   {
					/* Move low to normal */
					upriv->sched_node.next = upriv->low_sched_node.next;
					upriv->sched_node.prev = upriv->low_sched_node.prev;

					upriv->sched_node.next->prev = &upriv->sched_node;
					upriv->sched_node.prev->next = &upriv->sched_node;

					upriv->low_sched_node.next = upriv->low_sched_node.prev = &upriv->low_sched_node;

					/* And do yield */
					yield();
					
					continue;
			   }
		  }

		  s->prev->next = s->next;
		  s->next->prev = s->prev;

		  if (s != &f->sched_node)
		  {
			   fiber_t to = (fiber_t)((uintptr_t)s - (uintptr_t)&((fiber_t)0)->sched_node);
			   _fiber_switch(to);
		  }
		  break;
	 }
	 
	 upriv->in_sched = 0;
	 if (upriv->msg_incoming)
		  kmsg_clean_busy();
}
